function [optimization_results, power_allocation, fig_optimization] = capacity_optimization(EbN0_dB, optimization_types)
% 容量优化技术
% 输入参数：
%   EbN0_dB - 信噪比范围 (dB)
%   optimization_types - 优化类型 ('water_filling', 'power_allocation', 'adaptive_modulation', 'beamforming')
% 输出参数：
%   optimization_results - 优化结果
%   power_allocation - 功率分配结果
%   fig_optimization - 图形句柄

% 添加路径
addpath('../Common');
colors = color_definitions();

% 参数设置
bandwidth = 180e3; % 带宽 (Hz)
noise_power_dBm = -174 + 10*log10(bandwidth);
noise_power = 10^((noise_power_dBm-30)/10);
num_channels = 8; % 并行信道数量

% 初始化结果数组
num_snr = length(EbN0_dB);
num_types = length(optimization_types);
optimization_results = zeros(num_snr, num_types);
power_allocation = zeros(num_channels, num_types);

fprintf('容量优化技术...\n');

% 容量优化计算
for snr_idx = 1:num_snr
    snr_linear = 10^(EbN0_dB(snr_idx)/10);
    
    for type_idx = 1:num_types
        optimization_type = optimization_types{type_idx};
        
        switch lower(optimization_type)
            case 'water_filling'
                % 水填充算法
                % 生成随机信道增益 (dB)
                channel_gains_db = -10 + 5*randn(num_channels, 1);
                channel_gains = 10.^(channel_gains_db/10);
                
                % 水填充功率分配
                total_power = snr_linear * noise_power;
                power_allocated = water_filling(channel_gains, total_power, noise_power);
                
                % 计算总容量
                total_capacity = 0;
                for ch = 1:num_channels
                    if power_allocated(ch) > 0
                        channel_snr = power_allocated(ch) * channel_gains(ch) / noise_power;
                        total_capacity = total_capacity + bandwidth/num_channels * log2(1 + channel_snr);
                    end
                end
                optimization_results(snr_idx, type_idx) = total_capacity;
                
                % 保存功率分配 (中等SNR时)
                if snr_idx == round(num_snr/2)
                    power_allocation(:, type_idx) = power_allocated / total_power;
                end
                
            case 'power_allocation'
                % 功率分配优化
                % 多用户场景
                num_users = 4;
                user_channels = sqrt(0.5) * (randn(num_users, 1) + 1i * randn(num_users, 1));
                channel_gains = abs(user_channels).^2;
                
                % 最优功率分配 (基于信道质量)
                total_power = snr_linear * noise_power;
                power_allocated = optimal_power_allocation(channel_gains, total_power, noise_power);
                
                % 计算总容量
                total_capacity = 0;
                for user = 1:num_users
                    if power_allocated(user) > 0
                        user_snr = power_allocated(user) * channel_gains(user) / noise_power;
                        total_capacity = total_capacity + bandwidth/num_users * log2(1 + user_snr);
                    end
                end
                optimization_results(snr_idx, type_idx) = total_capacity;
                
            case 'adaptive_modulation'
                % 自适应调制
                % 多信道场景
                channel_gains = abs(sqrt(0.5) * (randn(num_channels, 1) + 1i * randn(num_channels, 1))).^2;
                
                % 自适应调制选择
                total_power = snr_linear * noise_power;
                power_per_channel = total_power / num_channels;
                
                total_capacity = 0;
                for ch = 1:num_channels
                    channel_snr = power_per_channel * channel_gains(ch) / noise_power;
                    modulation_order = select_modulation(channel_snr);
                    
                    % 计算该调制方式的容量
                    bits_per_symbol = log2(modulation_order);
                    total_capacity = total_capacity + bandwidth/num_channels * bits_per_symbol;
                end
                optimization_results(snr_idx, type_idx) = total_capacity;
                
            case 'beamforming'
                % 波束成形优化
                % MIMO场景
                Nt = 8; Nr = 4;
                num_streams = min(Nt, Nr);
                
                % 随机信道
                H = sqrt(0.5) * (randn(Nr, Nt) + 1i * randn(Nr, Nt));
                
                % SVD分解和波束成形
                [U, S, V] = svd(H);
                singular_values = diag(S);
                
                % 功率分配 (基于奇异值)
                total_power = snr_linear * noise_power;
                power_per_stream = total_power / num_streams;
                
                total_capacity = 0;
                for stream = 1:num_streams
                    if singular_values(stream) > 0
                        stream_snr = power_per_stream * singular_values(stream)^2 / noise_power;
                        total_capacity = total_capacity + bandwidth/num_streams * log2(1 + stream_snr);
                    end
                end
                optimization_results(snr_idx, type_idx) = total_capacity;
                
            otherwise
                % 默认等功率分配
                optimization_results(snr_idx, type_idx) = bandwidth * log2(1 + snr_linear);
        end
    end
end

%% 辅助函数
function power_allocated = water_filling(channel_gains, total_power, noise_power)
    % 水填充算法实现
    num_channels = length(channel_gains);
    power_allocated = zeros(num_channels, 1);
    
    % 排序信道增益
    [sorted_gains, sort_indices] = sort(channel_gains, 'descend');
    
    % 水填充
    water_level = 0;
    for k = 1:num_channels
        % 计算当前水位
        remaining_power = total_power;
        remaining_channels = k;
        
        if k > 1
            for i = 1:k-1
                remaining_power = remaining_power - max(0, water_level - noise_power/sorted_gains(i));
            end
        end
        
        if remaining_power > 0 && remaining_channels > 0
            water_level = (remaining_power + sum(noise_power./sorted_gains(1:k))) / k;
        end
    end
    
    % 分配功率
    for ch = 1:num_channels
        original_idx = sort_indices(ch);
        power_allocated(original_idx) = max(0, water_level - noise_power/channel_gains(original_idx));
    end
end

function power_allocated = optimal_power_allocation(channel_gains, total_power, noise_power)
    % 最优功率分配 (简化版)
    num_users = length(channel_gains);
    power_allocated = zeros(num_users, 1);
    
    % 基于信道质量的功率分配
    total_gain = sum(channel_gains);
    for user = 1:num_users
        power_allocated(user) = total_power * (channel_gains(user) / total_gain);
    end
end

function modulation_order = select_modulation(snr)
    % 自适应调制选择
    if snr < 5
        modulation_order = 2; % BPSK
    elseif snr < 10
        modulation_order = 4; % QPSK
    elseif snr < 15
        modulation_order = 16; % 16-QAM
    else
        modulation_order = 64; % 64-QAM
    end
end

%% 可视化结果
fig_optimization = figure('Name', '容量优化技术', 'Position', [100, 100, 1200, 800]);

% 子图1: 优化技术比较
subplot(2, 2, 1);
for type_idx = 1:num_types
    plot(EbN0_dB, optimization_results(:, type_idx)/1e6, 'LineWidth', 2, 'Color', colors(type_idx, :));
    hold on;
end
grid on;
xlabel('Eb/N0 (dB)');
ylabel('容量 (Mbps)');
title('容量优化技术比较');
legend(optimization_types, 'Location', 'NorthWest');

% 子图2: 功率分配比较
subplot(2, 2, 2);
bar(1:num_channels, power_allocation, 'grouped');
legend(optimization_types, 'Location', 'NorthEast');
xlabel('信道/用户索引');
ylabel('功率分配比例');
title('功率分配策略比较');

% 子图3: 优化增益
subplot(2, 2, 3);
equal_power_capacity = bandwidth * log2(1 + 10.^(EbN0_dB/10));
for type_idx = 1:num_types
    gain_db = 10*log10(optimization_results(:, type_idx) ./ equal_power_capacity');
    plot(EbN0_dB, gain_db, 'LineWidth', 2, 'Color', colors(type_idx, :));
    hold on;
end
grid on;
xlabel('Eb/N0 (dB)');
ylabel('优化增益 (dB)');
title('相对于等功率分配的增益');
legend(optimization_types, 'Location', 'NorthWest');

% 子图4: 容量效率
subplot(2, 2, 4);
efficiency = zeros(num_snr, num_types);
for snr_idx = 1:num_snr
    for type_idx = 1:num_types
        efficiency(snr_idx, type_idx) = optimization_results(snr_idx, type_idx) / bandwidth;
    end
end

for type_idx = 1:num_types
    plot(EbN0_dB, efficiency(:, type_idx), 'LineWidth', 2, 'Color', colors(type_idx, :));
    hold on;
end
grid on;
xlabel('Eb/N0 (dB)');
ylabel('频谱效率 (bps/Hz)');
title('频谱效率比较');

fprintf('容量优化技术完成！\n');

end